from nl2sql.model.llm import LLM
from nl2sql.prompt.pseudo_code_prompt import pseudo_code_prompt
from typing import Optional
import re


class PseudoCodeGenerator:

    def __init__(self, llm: Optional[LLM] = None):
        self.llm = llm
        self.prompt = pseudo_code_prompt

    @staticmethod
    def extract_pseudocode(text: str) -> str:
        """
        从完整文本中提取所有 LET 开头的伪代码语句块。
        返回格式化后的伪代码字符串。
        """
        pattern = r"(LET\s+[A-Z]+\s*=\s*.*?)(?=LET\s+[A-Z]+\s*=|$)"
        matches = re.findall(pattern, text, flags=re.DOTALL)
        cleaned = "\n".join([match.strip() for match in matches])
        return cleaned

    def generate(self, query: str, data_info: str) -> str:
        """
        根据 query 和 schema 生成伪代码表示
        """
        filled_prompt = self.prompt.format(query=query, data_info=data_info)
        response = self.llm.chat(filled_prompt)
        return self.extract_pseudocode(response)


# 示例用法
if __name__ == '__main__':
    from nl2sql.model.llm import LLM

    from nl2sql.datasource.excel_datasource import ExcelDatasource

    from nl2sql.model.llm import LLM
    from nl2sql.model.embedding_model import EmbeddingModel

    model_name = "qwen-max"
    api_key = "sk-34f5c792513c423f90b404d28b070f1f"
    base_url = "https://dashscope.aliyuncs.com/compatible-mode/v1"

    llm = LLM(model_name=model_name,
              api_key=api_key,
              base_url=base_url)

    embd = EmbeddingModel(api_key=api_key)

    excel_path = "D:/PythonProject_airline/nl2sql/examples/data"

    # 1. 初始化 Excel 数据源
    print("Loading Excel data source...")
    excel_ds = ExcelDatasource(source_path=excel_path, embedding_model=embd)

    # 2. 获取数据表信息
    data_info = excel_ds.get_data_info()
    query = "找出2024年销量大于100的产品在2025年的平均损耗率"
    pseudo_gen = PseudoCodeGenerator(llm)

    pseudo = pseudo_gen.generate(query=query, data_info=data_info)
    print(pseudo)
